"""
堆 (heap) 是一种经过排序的完全二叉树，
其中任一非叶子节点的值均不大于（或不小于）其左孩子和右孩子节点的值。
最大堆 根结点的键值是所有堆结点键值中最大者。
最小堆 根结点的键值是所有堆结点键值中最小者。
"""


class Heap(object):
    def __init__(self):
        self.data_list = []

    def get_parent_index(self, index):
        if index == 0 or index > len(self.data_list) - 1:
            return None
        else:
            """
            >> 位移运算符
            将 >> 左边的运算数二进位全部移动
            例：2 >> 1 的结果为 1
            该操作本质为 0000 0010 运算后为 0000 0001
            """
            return (index - 1) >> 1

    # 交换数组中的两个元素
    def swap(self, index_a, index_b):
        self.data_list[index_a], self.data_list[index_b] = self.data_list[index_b], self.data_list[index_a]

    def insert(self, data):
        """
        先把元素放到最后，然后从后往前依次堆化
        此处以最大堆为例，如果插入的元素比父节点大，则交换，直到最后
        """
        self.data_list.append(data)
        index = len(self.data_list) - 1
        parent = self.get_parent_index(index)
        # 循环，直到该元素成为堆顶，或小于父节点
        while parent is not None and self.data_list[parent] < self.data_list[index]:
            self.swap(parent, index)
            index = parent
            parent = self.get_parent_index(parent)

    def remove_max(self):
        # 删除堆顶元素，然后将最后一个元素放在堆顶，再从上往下依次堆化
        remove_data = self.data_list[0]
        self.data_list[0] = self.data_list[-1]
        del self.data_list[-1]

        self.heapify(0)
        return remove_data

    def heapify(self, index):
        # 从上往下堆化，从 index 开始堆化操作 (大顶堆)
        total_index = len(self.data_list) - 1
        while True:
            maxvalue_index = index
            if 2 * index + 1 <= total_index and self.data_list[2 * index + 1] > self.data_list[maxvalue_index]:
                maxvalue_index = 2 * index + 1
            if 2 * index + 2 <= total_index and self.data_list[2 * index + 2] > self.data_list[maxvalue_index]:
                maxvalue_index = 2 * index + 2
            if maxvalue_index == index:
                break

            self.swap(index, maxvalue_index)
            index = maxvalue_index

